# -*- coding:utf-8 -*-
# @FileName : app\models\conversation.py
# @Time     : 2025/11/28
# @Author   : 天空之城
"""对话模型文件

此模块定义对话和消息模型，使用SQLAlchemy进行ORM映射。
支持与agent_service.py的交互，包括工具调用记录、执行状态追踪等功能。
"""
from enum import Enum
from typing import Dict, Any, Optional, List
from datetime import datetime
from app.extensions import db


class ConversationMode(Enum):
    """对话模式枚举"""
    CHAT = "chat"
    TASK_RUNNER = "task_runner"
    SMART = "smart"


class ConversationStatus(Enum):
    """对话状态枚举"""
    IN_PROGRESS = "进行中"
    COMPLETED = "已完成"
    FAILED = "失败"


class MessageRole(Enum):
    """消息角色枚举"""
    USER = "user"
    ASSISTANT = "assistant"
    SYSTEM = "system"
    TOOL = "tool"


class MessageStatus(Enum):
    """消息状态枚举"""
    PENDING = "pending"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"


class Conversation(db.Model):
    """对话模型
    
    定义对话表结构，包含对话的基本信息和与用户的关联。
    支持与agent_service.py的交互，包括状态管理、工具调用追踪等功能。
    """
    __tablename__ = 'conversations'
    
    id = db.Column(db.Integer, primary_key=True, autoincrement=True)  # 对话ID
    user_id = db.Column(db.Integer, db.ForeignKey('users.id'), nullable=False)  # 所属用户ID
    title = db.Column(db.String(100), nullable=False, default='新对话')  # 对话标题
    mode = db.Column(db.Enum(ConversationMode), default=ConversationMode.CHAT)  # 对话模式
    status = db.Column(db.Enum(ConversationStatus), default=ConversationStatus.IN_PROGRESS)  # 对话状态
    execution_id = db.Column(db.String(100), nullable=True)  # 任务执行ID，用于追踪MCP任务
    meta_data = db.Column(db.JSON, nullable=True)  # JSON字段，存储执行上下文、工具调用统计等
    created_at = db.Column(db.DateTime, default=db.func.current_timestamp())  # 创建时间
    updated_at = db.Column(db.DateTime, default=db.func.current_timestamp(), onupdate=db.func.current_timestamp())  # 更新时间
    
    # 关系
    user = db.relationship('User', backref=db.backref('conversations', lazy=True))  # 关联的用户
    messages = db.relationship('Message', backref='conversation', lazy=True, order_by='Message.created_at')  # 关联的消息列表
    
    def __repr__(self):
        """返回对话对象的字符串表示"""
        return f'<Conversation {self.id} - {self.title}>'
    
    def update_status(self, status: ConversationStatus) -> None:
        """更新对话状态
        
        Args:
            status: 新的对话状态
        """
        self.status = status
        self.updated_at = datetime.utcnow()
        db.session.commit()
    
    def set_execution_id(self, execution_id: str) -> None:
        """设置任务执行ID
        
        Args:
            execution_id: 任务执行ID
        """
        self.execution_id = execution_id
        self.updated_at = datetime.utcnow()
        db.session.commit()
    
    def add_meta_data(self, key: str, value: Any) -> None:
        """添加元数据
        
        Args:
            key: 元数据键
            value: 元数据值
        """
        if self.meta_data is None:
            self.meta_data = {}
        self.meta_data[key] = value
        self.updated_at = datetime.utcnow()
        db.session.commit()
    
    def get_meta_data(self, key: str, default: Any = None) -> Any:
        """获取元数据
        
        Args:
            key: 元数据键
            default: 默认值
            
        Returns:
            元数据值或默认值
        """
        if self.meta_data is None:
            return default
        return self.meta_data.get(key, default)
    
    def get_tool_calls_summary(self) -> Dict[str, Any]:
        """获取工具调用统计摘要
        
        Returns:
            包含工具调用统计信息的字典
        """
        tool_calls = []
        for message in self.messages:
            if message.tool_calls:
                tool_calls.extend(message.tool_calls)
        
        return {
            'total_tool_calls': len(tool_calls),
            'successful_calls': len([call for call in tool_calls if call.get('status') == 'completed']),
            'failed_calls': len([call for call in tool_calls if call.get('status') == 'failed']),
            'tool_names': list(set([call.get('name', 'unknown') for call in tool_calls]))
        }


class Message(db.Model):
    """消息模型
    
    定义消息表结构，包含消息的基本信息和与对话的关联。
    支持与agent_service.py的交互，包括工具调用记录、执行状态追踪等功能。
    """
    __tablename__ = 'messages'
    
    id = db.Column(db.Integer, primary_key=True, autoincrement=True)  # 消息ID
    conversation_id = db.Column(db.Integer, db.ForeignKey('conversations.id'), nullable=False)  # 所属对话ID
    role = db.Column(db.String(20), nullable=False)  # 角色（user/assistant/system/tool）
    content = db.Column(db.Text, nullable=False)  # 文本内容
    image_path = db.Column(db.String(255), nullable=True)  # 截图文件路径
    tool_calls = db.Column(db.JSON, nullable=True)  # JSON字段，存储工具调用记录
    tool_results = db.Column(db.JSON, nullable=True)  # JSON字段，存储工具调用结果
    status = db.Column(db.Enum(MessageStatus), default=MessageStatus.PENDING)  # 消息状态
    meta_data = db.Column(db.JSON, nullable=True)  # JSON字段，存储耗时、Token消耗、执行上下文等
    created_at = db.Column(db.DateTime, default=db.func.current_timestamp())  # 创建时间
    
    def __repr__(self):
        """返回消息对象的字符串表示"""
        return f'<Message {self.id} - {self.role}>'
    
    def update_status(self, status: MessageStatus) -> None:
        """更新消息状态
        
        Args:
            status: 新的消息状态
        """
        self.status = status
        db.session.commit()
    
    def add_tool_call(self, tool_name: str, arguments: Dict[str, Any], 
                     call_id: Optional[str] = None) -> Dict[str, Any]:
        """添加工具调用记录
        
        Args:
            tool_name: 工具名称
            arguments: 工具调用参数
            call_id: 工具调用ID（可选）
            
        Returns:
            工具调用记录字典
        """
        if self.tool_calls is None:
            self.tool_calls = []
        
        tool_call = {
            'id': call_id or f'tool_call_{len(self.tool_calls) + 1}',
            'name': tool_name,
            'arguments': arguments,
            'status': 'pending',
            'created_at': datetime.utcnow().isoformat()
        }
        
        self.tool_calls.append(tool_call)
        db.session.commit()
        
        return tool_call
    
    def update_tool_call_result(self, call_id: str, result: Any, 
                              status: str = 'completed', error: Optional[str] = None) -> bool:
        """更新工具调用结果
        
        Args:
            call_id: 工具调用ID
            result: 工具调用结果
            status: 调用状态（completed/failed）
            error: 错误信息（可选）
            
        Returns:
            是否成功更新
        """
        if self.tool_calls is None:
            return False
        
        for tool_call in self.tool_calls:
            if tool_call.get('id') == call_id:
                tool_call['status'] = status
                tool_call['result'] = result
                tool_call['completed_at'] = datetime.utcnow().isoformat()
                if error:
                    tool_call['error'] = error
                
                # 更新工具调用结果列表
                if self.tool_results is None:
                    self.tool_results = []
                
                self.tool_results.append({
                    'call_id': call_id,
                    'tool_name': tool_call.get('name'),
                    'result': result,
                    'status': status,
                    'completed_at': datetime.utcnow().isoformat()
                })
                
                db.session.commit()
                return True
        
        return False
    
    def add_meta_data(self, key: str, value: Any) -> None:
        """添加元数据
        
        Args:
            key: 元数据键
            value: 元数据值
        """
        if self.meta_data is None:
            self.meta_data = {}
        self.meta_data[key] = value
        db.session.commit()
    
    def get_meta_data(self, key: str, default: Any = None) -> Any:
        """获取元数据
        
        Args:
            key: 元数据键
            default: 默认值
            
        Returns:
            元数据值或默认值
        """
        if self.meta_data is None:
            return default
        return self.meta_data.get(key, default)
    
    def get_tool_calls_by_status(self, status: str) -> List[Dict[str, Any]]:
        """根据状态获取工具调用列表
        
        Args:
            status: 工具调用状态
            
        Returns:
            符合条件的工具调用列表
        """
        if self.tool_calls is None:
            return []
        return [call for call in self.tool_calls if call.get('status') == status]
    
    def has_pending_tool_calls(self) -> bool:
        """检查是否有待处理的工具调用
        
        Returns:
            是否有待处理的工具调用
        """
        if self.tool_calls is None:
            return False
        return any(call.get('status') == 'pending' for call in self.tool_calls)
    
    def to_agent_format(self) -> Dict[str, Any]:
        """转换为agent_service.py使用的消息格式
        
        Returns:
            符合agent_service.py格式的消息字典
        """
        message_dict = {
            'role': self.role,
            'content': self.content
        }
        
        # 添加工具调用信息
        if self.tool_calls:
            message_dict['tool_calls'] = self.tool_calls
        
        # 添加工具调用结果信息
        if self.role == 'tool' and self.tool_results:
            message_dict['tool_call_id'] = self.tool_results[-1].get('call_id')
        
        # 添加元数据
        if self.meta_data:
            message_dict['meta'] = self.meta_data
        
        return message_dict